package com.es;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.Sequence;
import com.es.entity.Hotel;
import com.es.service.IHotelService;
import com.es.utils.ElasticUtil;
import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.dictionary.py.Pinyin;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.core.io.ClassPathResource;
import org.springframework.test.context.junit4.SpringRunner;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RunWith(SpringRunner.class)
@SpringBootTest
@Slf4j
public class ImportData {

    @Autowired
    private IHotelService hotelService;

    @Autowired
    private ElasticUtil elasticUtil;

    @Test
    public void importData() throws IOException {
        Sequence sequence = new Sequence();
        ClassPathResource resource = new ClassPathResource("db/hotel.txt");
        log.info("读取文件：{}", resource.getPath());
        List<String> list = FileUtil.readLines(resource.getFile(), CharsetUtil.charset("utf8"));
        List<Hotel> hotelList = new ArrayList<>();
        log.info("解析文件内容");
        list.stream().forEach(item -> {
            String[] strs = item.split("#");
            Hotel hotel = new Hotel();
            hotel.setId(sequence.nextId());
            hotel.setTitle(strs[1]);
            String businessDistrict = strs[2].replaceAll("\\[", "").replaceAll("\\]", "");
            hotel.setBusinessDistrict(StrUtil.join(",", businessDistrict.split("/")));
            hotel.setAddress(strs[3]);
            String[] location = strs[4].split(",");
            hotel.setLon(Double.parseDouble(location[0]));
            hotel.setLat(Double.parseDouble(location[1]));
            hotel.setCity(strs[5]);
            hotel.setPrice(Double.parseDouble(strs[6]));
            hotel.setStar(strs[7]);
            hotel.setFullRoom(Boolean.parseBoolean(strs[8]));
            hotel.setImpression(strs[9]);
            hotel.setFavourablePercent(Integer.parseInt(strs[10]));
            hotel.setPic(strs[11]);
            hotelList.add(hotel);
        });

        if (hotelList.size() > 0) {
            log.info("导入数据");
            hotelService.saveBatch(hotelList);
            log.info("导入成功");
        }
    }

    @Test
    public void importSuggestData() throws Exception {
        Sequence sequence = new Sequence();
        ClassPathResource resource = new ClassPathResource("db/suggest_words.csv");
        log.info("读取文件：{}", resource.getPath());
        List<String> list = FileUtil.readLines(resource.getFile(), CharsetUtil.charset("utf8"));
        log.info("解析文件内容");
        List<Map<String, Object>> mapList = new ArrayList<>();
        list.stream().forEach(item -> {
            String[] strs = item.split(",");
            mapList.add(processBy(sequence.nextId(), strs[0], Integer.parseInt(strs[1])));
        });

        if (mapList.size() > 0) {
            log.info("导入搜索建议");
            RestHighLevelClient highLevelClient = elasticUtil.getHighLevelClient();
            BulkRequest bulkRequest = new BulkRequest();
            mapList.stream().forEach(item -> {
                IndexRequest indexRequest = new IndexRequest("t_hotel_suggest");
                indexRequest.id(String.valueOf(item.get("id")));
                indexRequest.source(item);
                bulkRequest.add(indexRequest);
            });
            highLevelClient.bulk(bulkRequest, RequestOptions.DEFAULT);
            log.info("导入成功");
        }
    }

    public static Map<String, Object> processBy(long id, String text, int weight) {
        Map<String, Object> map = new HashMap<>();
        // 利用hanlp对中文进行处理
        List<Pinyin> pinyinList = HanLP.convertToPinyinList(text);
        // 提示词全拼
        StringBuffer fullPinyinBuffer = new StringBuffer();
        // 提示词拼音首字母
        StringBuffer headPinyinBuffer = new StringBuffer();
        pinyinList.stream().forEach(item -> {
            fullPinyinBuffer.append(item.getPinyinWithoutTone());
            headPinyinBuffer.append(item.getHead());
        });
        map.put("id", id);
        map.put("chinese", MapUtil.builder()
                .put("input", text)
                .put("weight", weight).build()
        );

        map.put("full_pinyin", MapUtil.builder()
                .put("input", fullPinyinBuffer.toString())
                .put("weight", weight).build()
        );

        map.put("head_pinyin", MapUtil.builder()
                .put("input", headPinyinBuffer.toString())
                .put("weight", weight).build()
        );
        return map;
    }
}
